package com.example.springbootdemo.websocket;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

@Component
public class ProgressWebSocketHandler extends TextWebSocketHandler {

    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.put(session.getId(), session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessions.remove(session.getId());
    }

    /**
     * 向指定会话发送进度更新
     * 
     * @param sessionId 会话ID
     * @param progress 进度值 (1-100)
     */
    public void sendProgressUpdate(String sessionId, int progress) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(String.valueOf(progress)));
            } catch (IOException e) {
                // 处理发送消息失败的情况
            }
        }
    }

    /**
     * 向所有连接的会话发送进度更新
     * 
     * @param progress 进度值 (1-100)
     */
    public void broadcastProgressUpdate(int progress) {
        sessions.values().forEach(session -> {
            if (session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(String.valueOf(progress)));
                } catch (IOException e) {
                    // 处理发送消息失败的情况
                }
            }
        });
    }
}